#Calculates Fast and Nathan's spin.
#Analysis by Alan Nathan, equations taken from Mike Fast:
#http://fastballs.wordpress.com/2007/09/07/magnus-papelbonus/
#
#The function returns only the newly calculated spin values, so you have to 
#cbind() it if you want to add it to an existing pitchfx() data frame.
#
#This function can be automatically called by get.pitches() using an option.
#
#UPDATE: apparently obsolete in 2009, as the gameday data now includes these variables

spin<-function(pitchfile) {
  spin.dir<-atan((pitchfile$az+32.174)/(pitchfile$ax))*(180/pi)
  spin.dir[which(pitchfile$ax<0)]<-spin.dir[which(pitchfile$ax<0)]+270
  spin.dir[which(pitchfile$ax>=0)]<-spin.dir[which(pitchfile$ax>=0)]+90
  
  spin.rate<-60*sqrt(I(pitchfile$ax^2)+I((pitchfile$az+32.174)^2))/
                    (abs(pitchfile$vy0)*0.121*0.00544*2*pi)
                    
  output<-cbind(spin.dir,spin.rate)
  colnames(output)<-c("spindir","spinrate")
  return(output)
}

#Use cluster analysis to classify pitches  
#
#Attempts to determine the number of different pitches using within-cluster sum
#of squares--increasing "tol" will cause it to distinguish more pitches.
#Uses at least speed and spin direction to find clusters, and optionally speed and break
#as well. Other variables can be specified.
#
#If naming=TRUE, the program then attempts to classify the pitch types it finds as 
#fastballs, curveballs, sliders or offspeed pitches, based on speed and spin.

classify.pitches <- function(pitcher,numpitches=NULL,
			    var1="start_speed",var2="spindir",var3="pfx_x",var4="pfx_z",var5=NULL,
			    tol=.885,naming=TRUE) {

  if(length(which(is.na(pitcher$start_speed)))>0 | sum(pitcher$pitch_type=="IN",na.rm=T)>0) {
    nodata<-which(is.na(pitcher$start_speed) | pitcher$pitch_type=="IN")
    reserved<-pitcher[nodata,]
    pitcher<-pitcher[-nodata,]
    reserved$spindir<-NA
    reserved$spinrate<-NA
  } else {
    reserved<-NULL
  }

  if ('throws' %in% names(pitcher)) hand <- ifelse(pitcher$throws[1]=="R",1,-1)
  if ('p_throws' %in% names(pitcher)) hand <- ifelse(pitcher$p_throws[1]=="R",1,-1)

  pitcher$spindir<-spin(pitcher)[,1]
  pitcher$spinrate<-spin(pitcher)[,2]

  if((mean(pitcher$spindir,na.rm=T)<210 & hand==1) | (mean(pitcher$spindir,na.rm=T)>150 & hand== -1)) {
    pitcher$spindir<-pitcher$spindir+60*hand
    pitcher$spindir[pitcher$spindir>360]<-pitcher$spindir[pitcher$spindir>360]-360
    pitcher$spindir[pitcher$spindir<0]<-pitcher$spindir[pitcher$spindir<360]+360
  }

  scaled.pitches<-data.frame(scale(pitcher[[var1]]),scale(pitcher[[var2]]))
  if(!is.null(var3)) scaled.pitches<-data.frame(scaled.pitches,scale(pitcher[[var3]]))
  if(!is.null(var4)) scaled.pitches<-data.frame(scaled.pitches,scale(pitcher[[var4]]))
  if(!is.null(var5)) scaled.pitches<-data.frame(scaled.pitches,scale(pitcher[[var5]]))

  if(is.null(numpitches)) {
      wss <- rep(NA,6)
      for(p in 2:7) wss[p-1] <-sum(kmeans(scaled.pitches,p,iter.max=50,nstart=25)$withinss)
      wss.diff <- wss
      for(p in 2:7) wss.diff[p-1] <- wss[p-1]/wss[1]
      drop <- wss.diff[1] - wss.diff[2]
      clusters <- 2
      while(drop>(1-tol) & clusters<length(wss.diff)) {
	clusters <- clusters + 1
	drop <- wss.diff[clusters-1]-wss.diff[clusters]
      }
  } else {
    clusters <- numpitches
  }

  pitch<-kmeans(scaled.pitches,clusters,iter.max=100,nstart=100)
  pitch.cluster<-pitch$cluster

  pitcher <- data.frame(pitcher, pitch$cluster)
 
  if(naming) {
    clusters <- unique(pitcher$pitch.cluster)
    speed.avg <- rep(NA,length(clusters))
    spin.avg <- rep(NA,length(clusters))
    for(i in 1:length(clusters)) {
      speed.avg[i] <- mean(pitcher[pitcher$pitch.cluster==clusters[i],'start_speed'],na.rm=T)
      spin.avg[i] <- mean(pitcher[pitcher$pitch.cluster==clusters[i],'spindir'],na.rm=T)
    }
    fb <- which(speed.avg==max(speed.avg,na.rm=T))
    pitcher$pitch.cluster[pitcher$pitch.cluster==clusters[fb]] <- "FB"
    clusters <- unique(pitcher$pitch.cluster)
    fb <- which(clusters=="FB")
    for(i in clusters[-fb]) {
	if(mean(pitcher$start_speed[pitcher$pitch.cluster=="FB"],na.rm=T)-
	   mean(pitcher$start_speed[pitcher$pitch.cluster==i],na.rm=T) < 3) {
	      pitcher$pitch.cluster[pitcher$pitch.cluster==i] <- paste("FB",i,collapse="")
	   }

	else if(hand*(mean(pitcher$spindir[pitcher$pitch.cluster=="FB"],na.rm=T)-
	   mean(pitcher$spindir[pitcher$pitch.cluster==i],na.rm=T)) > 20) {
        if((mean(pitcher$start_speed[pitcher$pitch.cluster=="FB"],na.rm=T)-
	         mean(pitcher$start_speed[pitcher$pitch.cluster==i],na.rm=T) < 8) | (hand*(mean(pitcher$spindir[pitcher$pitch.cluster=="FB"],na.rm=T)-
	   mean(pitcher$spindir[pitcher$pitch.cluster==i],na.rm=T)) < 100)) {
	           pitcher$pitch.cluster[pitcher$pitch.cluster==i] <- "SL"
           } else {
             pitcher$pitch.cluster[pitcher$pitch.cluster==i] <- "CU"
           }
	   }
	else if((hand*(mean(pitcher$spindir[pitcher$pitch.cluster=="FB"],na.rm=T) -
	   mean(pitcher$spindir[pitcher$pitch.cluster==i],na.rm=T)) < 10) & 
    (mean(pitcher$start_speed[pitcher$pitch.cluster=="FB"],na.rm=T)-
	   mean(pitcher$start_speed[pitcher$pitch.cluster==i],na.rm=T) >= 3)) {
	      pitcher$pitch.cluster[pitcher$pitch.cluster==i] <- "OS"
	   }	
    }
  }
  
  if(!is.null(reserved)) {
    reserved$pitch.cluster<-NA
    pitcher<-rbind(pitcher,reserved)
  }
  
  pitcher$pitch.cluster <- factor(pitcher$pitch.cluster)
  pitcher$spindir<-spin(pitcher)[,1]
  class(pitcher) <- c("pitchfx","data.frame")

  return(pitcher)

}

#Model the probability that a pitch will be called a strike
#Uses a logistic regression model with two quadratic terms, using balls not swung at as inputs
#Uses that model to predict how swung-at pitches would have been called if they had been taken.
strikeprob <- function(test,random.prob=FALSE,return.model=FALSE) {
  require(arm)
  
  m1<-bayesglm(calledstrike~px+I(px^2)+pz+I(pz^2),data=test,family=binomial(link="logit"))

  if(return.model) return(m1)
    
  pred.strike<-function(h,v) {
    if(random.prob) { 
      sim1<-sim(m1,n.sims=length(h))
      output<-invlogit(rowSums(sim1$beta * cbind(1,h,h*h,v,v*v)) + rnorm(length(h),0,sim1$sigma))
      return(output)
    }
    if(!random.prob) return(invlogit(cbind(1,h,h*h,v,v*v) %*% c(coef(m1)[1],coef(m1)[2],coef(m1)[3],coef(m1)[4],coef(m1)[5]))) 
  }
  
  pred.strike.z<-function(h,v) {
    if(random.prob) { 
      sim1<-sim(m1,n.sims=length(h))
      output<-rowSums(sim1$beta * cbind(1,h,h*h,v,v*v)) + rnorm(length(h),0,sim1$sigma)
      return(output)
    }
    if(!random.prob) return(cbind(1,h,h*h,v,v*v) %*% c(coef(m1)[1],coef(m1)[2],coef(m1)[3],coef(m1)[4],coef(m1)[5])) 
  }
  
  test$strikeprob<-pred.strike(test$px,test$pz)
  test$strikeprob.z<-pred.strike.z(test$px,test$pz)
  
  #test$strikeprob<-NA
  #for(i in 1:dim(test)[1]) {
  #  test$strikeprob[i]<-pred.strike(test$px[i],test$pz[i])
  #}
  
  subset<-which(!is.na(test$strikeprob))
  test$predstrike<-NA
  test[subset,]$predstrike<-rbinom(dim(test[subset,])[1],1,test[subset,]$strikeprob)
  
  output<-cbind(test$strikeprob,test$predstrike,test$strikeprob.z)
  colnames(output)<-c("Strike prob.","Strike prediction","Strike propensity (latent)")
  return(output)

}

#Models the probability that a player will swing at a pitch, as a function of the probability that the pitch will be called a strike. 
#Uses Gelman et al's bayesglm() with the default prior, in order to avoid separation problems. Returns a bayesglm object.
#By default, just regresses whether the batter swung on strike probability and number of strikes, but different models can be specified.
selectivity.curve<-function(player,model=formula(swung~probs+factor(s)),random.prob=FALSE) {
  require(arm)
  player<-player[which(!is.na(player$outcome)),]
  predictions<-strikeprob(player,random.prob=random.prob)
  player$probs<-predictions[,1]
  player$z <- predictions[,3]
  player$swung<-ifelse(player$outcome %in% c("Foul","In play, out","Reached base","Swinging strike"),1,0)
  selectivity<-bayesglm(model,data=player,family=binomial(link="logit"))
  return(selectivity)
}

#Model pitch trajectory based on the kinematic equation.
#Based on Dan Fox's description:
#http://www.baseballprospectus.com/article.php?articleid=6569
pitchloc <- function(t,x0,ax,vx0,y0,ay,vy0,z0,az,vz0) {
  x <- x0 + vx0*t + 0.5*ax*I(t^2)
  y <- y0 + vy0*t + 0.5*ay*I(t^2)
  z <- z0 + vz0*t + 0.5*az*I(t^2)  
  if(length(t)==1) {
    loc<-c(x,y,z)
  } else {
    loc <- cbind(x,y,z)
  }
  return(loc)
}

#Coordinate shift for batter's-box viewpoint
#Implemented from Mike fast's explanation:
#http://fastballs.wordpress.com/2008/06/15/view-from-the-batters-box/
shift <- function(x,y,z,x0,offset=2.8) {
  m <- sqrt(y^2 + (x + offset)^2)
  alpha <- atan( 55/(x0 + offset) )
  theta <- alpha - atan( y/(x+offset) )
  y1 <- m*cos(theta)
  x1 <- m*sin(theta)
  out<-cbind(x1,y1,z)
  colnames(out)<-c("x","y","z")
  out<-data.frame(out)
  return(out)
}

#Generates a model of a pitch's trajectory, averaging over all the pitches in the
#input dataset. Works best with data containing only a single pitch type (using
#either the gameday identification, or better, clusters from pitch.classify().)
pitch.trajectory <- function(pitches,interval=.01) {
  pitches<-pitches[complete.cases(pitches$x0,pitches$ax,pitches$vx0,pitches$y0,pitches$ay,pitches$vy0,pitches$z0,pitches$az,pitches$vz0),]
  x0 <- mean(pitches$x0,na.rm=T)
  ax <- mean(pitches$ax,na.rm=T)
  vx0 <- mean(pitches$vx0,na.rm=T)
  y0 <- mean(pitches$y0,na.rm=T)
  ay <- mean(pitches$ay,na.rm=T)
  vy0 <- mean(pitches$vy0,na.rm=T)
  z0 <- mean(pitches$z0,na.rm=T)
  az <- mean(pitches$az,na.rm=T)
  vz0 <- mean(pitches$vz0,na.rm=T)

  #Use quadratic formula to figure out when the pitch crosses the plate
  cross.plate <- (-1*vy0 - sqrt(I(vy0^2)-2*y0*ay))/ay

  tracking<-t(sapply(seq(0,cross.plate,interval),pitchloc,x0=x0,ax=ax,vx0=vx0,y0=y0,ay=ay,vy0=vy0,z0=z0,az=az,vz0=vz0))
  colnames(tracking)<-c("x","y","z")
  tracking<-data.frame(tracking)

  return(tracking)    
}

#Analyzes batter behavior in different areas of the strike zone.
#
#Returns a list. The first element is an array of six matrices: 
#swing rates, contact rates, in play rates, hit rates,
#number of pitches seen, and number of swings for each region.
#
#The number of zones to subdivide the strike zone into is given in "breaks". The
#"buffer" parameter can be used to adjust the boundaries of the area to be
#analyzed. The "center" option adjusts the hitting zone location according to
#the average location of pitches in the data.
#
#This function works best with a dataset containing a large number of pitches.

hit.zones <- function(pitches,breaks=4,buffer=c(0.2,0.2,0.2,0.2),center=T) {

  pitches<-pitches[(!is.na(pitches$px) & !is.na(pitches$pz)),]
  swung <- ifelse(pitches$outcome %in% c("Foul","In play, out","Reached base","Swinging strike"),1,0)
  x.all<-pitches$px
  y.all<-pitches$pz
  pitches<-pitches[pitches$outcome %in% c("Foul","In play, out","Reached base","Swinging strike"),]

  x<-pitches$px
  y<-pitches$pz
  
  contact <- ifelse(pitches$outcome %in% c("Foul","In play, out","Reached base"),1,0)
  inplay <- ifelse(pitches$outcome %in% c("In play, out","Reached base"),1,0)
  hit <- ifelse(pitches$hit_event %in% c("Double","Home Run","Line Out","Single","Triple"),1,0)

  xoffset <- yoffset <- 0
  if(center) {
    xoffset<-mean(x[x > -2 & x < 2])
    yoffset<-mean(y[y > 0.5 & y < 4.5])-
        (mean(pitches$sz_top,na.rm=T)-
          0.5*(mean(pitches$sz_top,na.rm=T)-mean(pitches$sz_bot,na.rm=T)))
  }
  
  sz.top <- mean(pitches$sz_top,na.rm=T)+buffer[1]+yoffset
  sz.bot <- mean(pitches$sz_bot,na.rm=T)-buffer[3]+yoffset
  sz.left <- -1 - buffer[2]+xoffset
  sz.right <- 1 + buffer[4]+xoffset

  v.bounds <- c(sz.top,sz.top-((1:(breaks-1))/breaks)*(sz.top-sz.bot),sz.bot)
  h.bounds <- c(sz.right,sz.left-(((breaks-1):1)/breaks)*(sz.left-sz.right),sz.left)

  result <- array(NA,c(breaks,breaks,6))
  for(i in 1:breaks) {
    for(j in 1:breaks) {
      zone0 <- which(y.all<v.bounds[i] & y.all>v.bounds[i+1]
              & x.all<h.bounds[j] & x.all>h.bounds[j+1])
      zone <- which(y<v.bounds[i] & y>v.bounds[i+1]
                  & x<h.bounds[j] & x>h.bounds[j+1])
      swung.temp <- swung[zone0]
      contact.temp <- contact[zone]
      inplay.temp <- inplay[zone]
      hit.temp <- hit[zone]

      result[i,j,1]<-mean(swung.temp,na.rm=T)
      result[i,j,2]<-mean(contact.temp,na.rm=T)
      result[i,j,3]<-mean(inplay.temp,na.rm=T)
      result[i,j,4]<-mean(hit.temp,na.rm=T)
      result[i,j,5]<-length(contact.temp)
      result[i,j,6]<-length(swung.temp)
    }
  }

  totals<-c(mean(swung,na.rm=T),mean(contact,na.rm=T),mean(inplay,na.rm=T),mean(hit,na.rm=T))
  names(totals)<-c("Swing","Contact","In play","Hit")
  results <- list(zones=result,totals=totals,buffer=buffer,offset=c(xoffset,yoffset))
  return(results)
}
